package com.xiaoc.bulldozer.springboot.config.wrapper;

import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.*;

/**
 * @author wj on 2019/5/15.
 */
public class XiaocResponseWrapper extends HttpServletResponseWrapper {

    private ByteArrayOutputStream outputStream  = new ByteArrayOutputStream();
    private PrintWriter printWriter;

    /**
     * Constructs a response adaptor wrapping the given response.
     *
     * @param response The response to be wrapped
     * @throws IllegalArgumentException if the response is null
     */
    public XiaocResponseWrapper(HttpServletResponse response) {
        super(response);
    }

    @Override
    public PrintWriter getWriter() throws IOException {
        try{
            printWriter = new PrintWriter(new OutputStreamWriter(outputStream, "utf-8"));
        } catch(UnsupportedEncodingException e) {
            e.printStackTrace();
        }
        return printWriter;
    }


    public byte[] getBytes() {
        if(null != printWriter) {
            printWriter.close();
            return outputStream.toByteArray();
        }

        if(null != outputStream) {
            try {
                outputStream.flush();
            } catch(IOException e) {
                e.printStackTrace();
            }
        }
        return outputStream.toByteArray();
    }

}
